
from mootdx.quotes import Quotes
import time
import pandas as pd


def getBars(symbol='', period='', count=1) -> pd.DataFrame:
    """封装实时行情Bars接口调用
    入参：arg[0]: symbol, arg[1]: period(5/15/30/60/120/240), arg[2]: count
    返回：open,close,high,low,vol,amount,year,month,day,hour,minute,datetime
    """

    # TDX接口实现
    global client
    client = Quotes.factory(market='std', multithread=True)
    CLIENT_QUERY_PERIOD = 0.1  # 访问服务器流控：单位秒

    # tdx查询实时K线数据的frequency参数与周期对应表, 120不支持
    TDX_FREQS = {'1': 8, '5': 0, '15': 1, '30': 2, '60': 3, '120': -1, '240': 4}

    if period == '120':  # 120分钟周期数据从60分钟数据提取
        rd = client.bars(symbol=symbol, frequency=TDX_FREQS['60'], offset=count*2)
        for i, data in rd.iterrows():
            if i % 2 == 0:
                rd.drop(index=i, inplace=True)
        rd.reset_index(drop=True, inplace=True)
    else:
        rd = client.bars(symbol=symbol, frequency=TDX_FREQS[period], offset=count)

    time.sleep(CLIENT_QUERY_PERIOD)
    return rd


if __name__ == "__main__":
    print(getBars('159995', '15', 10))
